import random

from common.Common import Point
import numpy as np
import math

from common.Constant import Constant

# todo 随机生成的task对结果影响较大

# 边缘服务器
class EdgeServer:
    def __init__(self, id: int, location: Point, bandwidth: float = 1, capacitance: float = 1e-29):
        # CPU, memory, storage and bandwidth
        # we assume that there are sufficient memory, storage and network capacity in each node.
        # 每个容器分配的计算资源 取决于容器数量 （也可以用task量加权分配）

        self.id = id  # 边缘服务器id, 作为环境中边缘服务器列表下标
        self.location = location
        # 服务器到车辆的总带宽 上行 和 下行
        self.uplink = 10.0  # Mhz 上行链路带宽
        self.bandwidth = bandwidth  # 两个边缘服务器之间的带宽 Gb/s
        # 服务器之间通信的总带宽 上行 和 下行
        # todo 服务器覆盖范围

        self.no = 1e-16
        self.user_index = []  # 服务用户的下标 用来检索container task
        self.capacitance = capacitance  # effective switch capacitance

        # F 边缘服务器总计算能力
        self.F = 20.0  # 边缘服务器总计算能力 GHz
        self.container_num = 0  # 部署在es上的container 数量
        self.tran_power = 20.0  # w 传输能耗

    def reset(self):
        self.container_num = 0
        self.user_index = []


# 容器（就是Service 在mec中迁移）
# 因为一个用户被一个服务的服务，所以Container与User是一对一关系
class Container:
    def __init__(self, id: int, container_size: float = 1.0):
        self.id = id
        self.data_size = container_size  # Mb


# 移动用户（包括车辆等）
class User:
    # todo 写一个方法可以把user的状态转换为numpy
    # todo 写一个方法返回每个下表对应的是什么含义
    # todo 每个用户都有卸载数据大小 task
    def __init__(self, user_id: int, location: list[Point], container: Container, direction: int, step_distance):
        self.user_id = user_id
        self.location = location  # 用户坐标
        self.container = container  # 考虑要不要在这里
        self.direction = direction  # 0是向右，从小到大。 1是向左 从大到小
        self.step_distance = step_distance  # 模拟speed todo 改为可变
        self.pn = 1  # w  # Transmit power of each gNB
        self.task = None

    def __str__(self) -> str:
        return f"user_id:{self.user_id},location:{self.location},container:{self.container}" \
               f",direction:{self.direction}."



# MEC网络的总控制器
class SDNController:
    # 每个MEC服务器都扮演一个代理的角色，它独立地做出服务放置和迁移决策
    # 一旦每个MEC代理做出决策，它就会将其传达给SDN控制器，该控制器扮演中央代理的角色，以协调所有MEC服务器的决策。

    def check_action(self, action):
        '''
        检查action合法性
        :return:
        '''
        user_num = len(action[0])
        for i in range(user_num):
            tmp = sum([row[i] for row in action])
            if tmp != 1:
                return False
        return True

    def make_global_action(self, args):
        '''
        可变参数，每个agent传递的action拼合成global_action
        :param args:
        :return:
        '''
        return []


# 用户传输给mec需要被Container处理的数据
# 每个用户每个时隙生成0.2M大小的计算请求
class Task:
    def __init__(self, random_rate: float = 0.5, f0: float = 50.0):
        self.data_size = np.random.randint(1, random_rate) / 10 * 1024 * 1024 * 8  # bit 迁移服务的数据大小
        self.f0 = f0  # 每bit需要的cpu周期数 Mhz



# MEC环境
class MECEnv:
    # state (current user position[], 速度[], 方向[], 用户在和哪台服务器通信, es[])

    # 问题是es[]如何加入到 state里面表示 : 一行 是一个用户的信息，和所有服务器的信息 参考multiagentmecoffload项目
    # 多关注其他论文里面state是怎样设计的

    def __init__(self, edge_servers: list[EdgeServer], users: list[User], containers: [Container],
                 time_slot_total: int, tasks_random_rate: list[int]):
        # ES

        self.edge_servers = edge_servers
        self.es_num = len(self.edge_servers)
        # User
        self.users = users
        self.user_num = len(self.users)
        # Container len(users)=len(containers)
        self.containers = containers
        self.action_dim = self.es_num
        self.state_dim = 7
        self.time_slot_now = 0
        self.time_slot_total = time_slot_total  # 待设定
        self.sdn = SDNController()

        self.state = []
        self.reward = 0
        self.w1 = 10.0  # 延迟占比
        self.w2 = 1.0  # 能耗占比
        self.tasks_random_rate = tasks_random_rate # 根据概率生成task

    def __reward(self):
        pass

    def __get_communication_delay(self, user: User, server: EdgeServer, task: Task):
        '''
        用户上传task的通信延迟
        :param user: 用户
        :param server: 边缘服务器
        :param task: offload的任务
        :return: 数值
        '''
        # num = 1 if server.container_num<=0 else server.container_num
        # tmp_rn = server.uplink / server.container_num  # W / K 速率公式的一部分
        tmp_rn = server.uplink / self.user_num
        distance = Point.distance(user.location[self.time_slot_now], server.location)
        t = 1e-6
        h = t * pow(distance, -2)  # 无线增益
        # h = pow(distance, -3)

        # 数据传输率
        data_trans_rate = tmp_rn * math.log2(1 + user.pn * h / server.no)  # 计算速率 (1 + transmission power *无线信道增益 /  信噪比)
        communication_delay = task.data_size / data_trans_rate
        return communication_delay

    def __get_computing_delay(self, task: Task, server: EdgeServer):
        '''
        服务的计算延迟
        :param task:
        :param server:
        :return:
        '''
        f_com = (server.F * 1000 / server.container_num) * 1000  # hz 分配给改任务的计算能力 计算能力均分
        computing_delay = task.data_size * task.f0 / f_com
        return computing_delay

    def __get_migration_delay(self, container: Container, server_from: EdgeServer, server_to: EdgeServer):
        if server_from.id == server_to.id:
            # 没有发生迁移
            return 0

        migration_delay = container.data_size / (server_to.bandwidth * 1024)
        return migration_delay
        # 论文中没给公式
        # todo 这个不会因为就离变高嘛
        # 能量消耗，

    def __get_computing_cost(self, task: Task, server: EdgeServer):
        # 边缘服务器计算任务产生的能耗
        f_com = (server.F * 1000 / server.container_num) * 1000  # hz 分配给该任务的计算能力 计算能力均分
        computing_cost = server.capacitance * task.data_size * task.f0 * f_com ** 2
        return computing_cost

    def __get_migration_cost(self, container: Container, server_from: EdgeServer, server_to: EdgeServer):
        return server_to.tran_power * self.__get_migration_delay(container, server_from, server_to)

    def __read_state_from_file(self, ):
        '''
        从文件里获取state的用户位置，方向，速度等
        :param file:
        :return:
        '''
        return []

    # reward 应该是执行了这个action 之后的奖励
    def step(self, action):
        # 距离是根据哪个state算的

        if self.time_slot_now >= self.time_slot_total - 1:
            return self.state, self.reward, True, {}

        # 计算state
        self.time_slot_now += 1
        old_state = self.state
        # 获取old_state中 container都部署在哪个服务器上
        old_allocation = []
        for single_state in old_state:
            old_allocation.append(int(single_state[Constant.service_server_index]))

        next_state = []
        for server in self.edge_servers:
            server.reset()

        for i in range(self.user_num):
            user = self.users[i]
            user_position = user.location[self.time_slot_now].get_list()
            user_direction = user.direction
            user_speed = user.step_distance
            container_size = user.container.data_size
            task = Task(self.tasks_random_rate[i])
            user.task = task
            task_size = task.data_size
            service_server = action[i]
            self.edge_servers[service_server].container_num += 1
            self.edge_servers[service_server].user_index.append(i)
            # server_x = []
            # server_y = []
            # for i in range(self.es_num):
            #     es = self.edge_servers[i]
            #     server_x.append(es.location.point_x)
            #     server_y.append(es.location.point_y)
            #     # todo 链路变化
            single_state = np.array([user_direction, user_speed, container_size, task_size, service_server])
            # single_state = np.hstack((user_position, single_state, server_x, server_y))
            single_state = np.hstack((user_position, single_state))
            # todo containner 数目
            next_state.append(single_state)
        self.state = np.array(next_state, dtype=np.float32)

        next_allocation = []
        for single_state in self.state:
            next_allocation.append(int(single_state[Constant.service_server_index]))

        # delay
        commu_dely = 0
        compu_delay = 0
        mig_delay = 0
        compu_cost = 0
        mig_cost = 0

        for i in range(self.user_num):
            user = self.users[i]
            server_to = self.edge_servers[next_allocation[i]]
            server_from = self.edge_servers[old_allocation[i]]  # todo 这里其实是心状态下的server container serviceuser 都不准确
            task = user.task
            container = self.containers[i]
            commu_dely += self.__get_communication_delay(user, server_to, task)
            compu_delay += self.__get_computing_delay(task, server_to)
            mig_delay += self.__get_migration_delay(container, server_from, server_to)  # 太小了
            compu_cost += self.__get_computing_cost(task, server_to)
            mig_cost += self.__get_migration_cost(container, server_from, server_to)  # 太小了

        # reward = -1 / self.user_num * (self.w1 * (commu_dely + compu_delay + mig_delay) + self.w2 * (
        #         compu_cost + mig_cost))
        reward = -commu_dely / self.user_num
        self.reward = reward
        return self.state, self.reward, False, {}

    def reset(self):
        self.time_slot_now = 0
        # user_x, user_y, user_direction, user_speed, container_size, task_size, service_server, server_x[], server_y[], container_num
        state = []
        for i in range(self.es_num):
            self.edge_servers[i].reset()

        for i in range(self.user_num):
            user = self.users[i]
            user_position = user.location[self.time_slot_now].get_list()
            user_direction = user.direction
            user_speed = user.step_distance
            container_size = user.container.data_size
            task = Task(self.tasks_random_rate[i])
            user.task = task
            task_size = task.data_size
            service_server = random.randint(0, self.es_num - 1)
            self.edge_servers[service_server].container_num += 1
            self.edge_servers[service_server].user_index.append(i)
            # server_x = []
            # server_y = []
            # for i in range(self.es_num):
            #     es = self.edge_servers[i]
            #     server_x.append(es.location.point_x)
            #     server_y.append(es.location.point_y)
            #     # todo 链路变化
            single_state = np.array([user_direction, user_speed, container_size, task_size, service_server])
            # single_state = np.hstack((user_position, single_state, server_x, server_y))
            single_state = np.hstack((user_position, single_state))
            # todo containner 数目
            state.append(single_state)
        self.state = np.array(state, dtype=np.float32)
        return self.state

    @staticmethod
    def env(user_num: int = 2, server_num: int = 5):
        # 长 5000米 宽6米
        step_distance = 5  # m
        es_distance = 1000
        total_distance = server_num * es_distance
        total_step = int(total_distance / step_distance)

        # 获取user的所有轨迹
        # 生成users
        # 生成servers
        ESs = []
        for i in range(server_num):
            ESs.append(EdgeServer(i, Point((i) * es_distance, 3)))

        containers = []
        users = []
        for i in range(user_num):
            # todo container的数据大小可变
            container = Container(i, 0.5*i+1)
            containers.append(container)
            if i % 2 == 0:
                # for i in range(int(total_step)):
                #     points.append(Point(i * step_distance, 2))
                points = [Point(i * step_distance, 2) for i in range(int(total_step))]
                users.append(User(i, points, container, 0, step_distance))
            else:
                points = [Point(total_distance - i * step_distance, 4) for i in range(int(total_step))]
                users.append(User(i, points, container, 1, step_distance))
        # todo user 的speed可变 或者不同
        # todo 写一个seed做环境的随机状态生成
        tasks_random_rate = [5+i for i in range(user_num)]
        return MECEnv(ESs, users, containers, total_step, tasks_random_rate)


if __name__ == "__main__":
    env = MECEnv.env()
    print('Suc')
    print(env.reset())
    state, reward, done, info = env.step([0, 1])
    state, reward, done, info = env.step([0, 2])
    state, reward, done, info = env.step([0, 3])
    state, reward, done, info = env.step([0, 4])

    for i in range(100):
        print(np.random.randint(1,10))